﻿# coding=utf-8
# 2023/4/27: 首个记录，基于旧2020/8/5版本
# 2023/5/5: 枚举基于Enum
# 2023/10/10: 补充部分英文注释
# 2025/7/18: 修正significant赋值
# 2025/7/21: 使用类型注解。优化to_general_sample。修正orientation插值

import bi_common as bi
from math import pi, sin, cos, atan2
from enum import Enum
from typing import List, Any


class DynamicObjectType(Enum):
    UNKNOWN = 0  # 未知
    CAR = 1  # 轿车
    VAN = 2  # 面包车
    TRUCK = 3  # 货车
    TRAILER = 4  # 拖车
    SEMITRAILER = 5  # 货柜车
    BUS = 6  # 大巴
    MOTORBIKE = 7  # 摩托车
    BICYCLE = 8  # 自行车
    TRAIN = 9  # 火车
    TRAM = 10  # 有轨电车
    PEDESTRIAN = 11  # 行人
    WHEELCHAIR = 12  # 轮椅
    ANIMAL = 13  # 动物


class DynamicObject:
    def __init__(self):
        self.id: int = 0  # ID of the object, 0 as subject, 1~ as other objects / 目标物ID，0表示为主目标物，1~表示一般目标物
        self.type: DynamicObjectType = DynamicObjectType.UNKNOWN  # Classification of the object / 目标物类型
        self.center_longitude: float = 0.0  # [°E] Longitude of the object's box center / 目标物中心的经度
        self.center_latitude: float = 0.0  # [°N] Latitude of the object's box center / 目标物中心的纬度，北为正方向
        self.center_altitude: float | None = None  # [m] Altitude of the object's box center / 目标物中心的海拔
        self.orientation: float | None = None  # [deg] Heading of the object's box, 0 as north, CCW is positive / 目标物朝向，北为0，逆时针为正方向
        self.width: float | None = None  # [m] Width of the object's box / 目标物宽度
        self.length: float | None = None  # [m] Length of the object's box / 目标物长度
        self.height: float | None = None  # [m] Height of the object's box / 目标物高度
        self.speed: float | None = None  # [KPH] Speed of the object's box along the heading / 目标物速度
        self.accel_x: float | None = None  # [m/s²] Acceleration of the object's box along the heading / 目标物纵向加速度
        self.curvature: float | None = None  # [1/m] Turn curvature of the object's rear axis center / 目标物转弯曲率
        self.yaw_rate: float | None = None  # [deg/s] Yaw rate of the object's box / 目标物横摆角速度
        self.pitch: float | None = None  # [deg] Pitch angle of the object's box, heading down is positive / 目标物俯仰角，车头朝下为正
        self.roll: float | None = None  # [deg] Roll angle of the object's box, right side heading down is positive / 目标物横滚角，右侧朝下为正
        self.rear_axis_longitudinal_offset: float | None = None  # [m] Longitudinal distance between the object's box center and rear axis center, always positive / 目标物自身坐标系下，后轴中心与物体中心的纵向距离，恒为正
        self.rear_axis_height_offset: float | None = None  # [m] Height difference of the object's box center and rear axis center, always positive / 目标物自身坐标系下，后轴中心与物体中心的高度差，恒为正
        self.color_valid: bool = False  # Whether the color is valid / 目标物颜色是否有效
        self.color_r: int = 0  # Red component of the n-th object's color / 目标物颜色R分量
        self.color_g: int = 0  # Green component of the n-th object's color / 目标物颜色G分量
        self.color_b: int = 0  # Blue component of the n-th object's color / 目标物颜色B分量


class GlobalObjectsSample:
    def __init__(self):
        self.time: float = 0.0  # [s] Time offset in session / Session内的相对时间
        self.objects: List[DynamicObject] = []  # List of objects / 目标物列表

    # Convert to general sample for output / 转通用样本，用于样本输出
    def to_general_sample(self) -> Any:
        output = bi.agency.create_general_sample()
        output.protocol = 'global-objects-sample-v2'
        output.time = self.time
        
        total_length = int(1 + len(self.objects) * 26)
        output.significant = total_length
        
        values: List[float | str | None] = [None] * total_length
        values[0] = len(self.objects)
        
        for i, obj in enumerate(self.objects):
            base_index = 1 + i * 26
            values[base_index] = obj.id
            values[base_index + 1] = obj.type.value
            values[base_index + 2] = obj.center_longitude
            values[base_index + 3] = obj.center_latitude
            values[base_index + 4] = obj.center_altitude
            values[base_index + 5] = obj.orientation
            values[base_index + 6] = obj.width
            values[base_index + 7] = obj.length
            values[base_index + 8] = obj.height
            values[base_index + 9] = obj.speed
            values[base_index + 10] = obj.accel_x
            values[base_index + 11] = obj.curvature
            values[base_index + 12] = obj.yaw_rate
            values[base_index + 13] = obj.pitch
            values[base_index + 14] = obj.roll
            values[base_index + 15] = obj.rear_axis_longitudinal_offset
            values[base_index + 16] = obj.rear_axis_height_offset
            values[base_index + 17] = obj.color_r if obj.color_valid else None
            values[base_index + 18] = obj.color_g if obj.color_valid else None
            values[base_index + 19] = obj.color_b if obj.color_valid else None

        output.values = values
        return output


def _conv_global_objects_sample_v1(gs: Any) -> GlobalObjectsSample | None:
    values_count = len(gs.values)
    if values_count < 1:
        return None
    object_count = int(gs.values[0]) if gs.values[0] is not None else 0
    if values_count != 1 + 10 * object_count:
        return None
    output = GlobalObjectsSample()
    output.time = gs.time
    for i in range(0, object_count):
        obj = DynamicObject()
        b = 1 + 10 * i
        obj.id = int(gs.values[b] if gs.values[b] is not None else 0)
        obj.type = DynamicObjectType(int(gs.values[b + 1])) if gs.values[b + 1] is not None else DynamicObjectType.UNKNOWN
        obj.center_longitude = float(gs.values[b + 2]) if gs.values[b + 2] is not None else 0.0
        obj.center_latitude = float(gs.values[b + 3]) if gs.values[b + 3] is not None else 0.0
        obj.center_altitude = float(gs.values[b + 4]) if gs.values[b + 4] is not None else None
        obj.orientation = float(gs.values[b + 5]) if gs.values[b + 5] is not None else None
        obj.width = float(gs.values[b + 6]) if gs.values[b + 6] is not None else None
        obj.length = float(gs.values[b + 7]) if gs.values[b + 7] is not None else None
        obj.height = float(gs.values[b + 8]) if gs.values[b + 8] is not None else None
        obj.speed = float(gs.values[b + 9]) if gs.values[b + 9] is not None else None
        output.objects.append(obj)
    return output


def _conv_global_objects_sample_v2(gs: Any) -> GlobalObjectsSample | None:
    values_count = len(gs.values)
    if values_count < 1:
        return None
    object_count = int(gs.values[0]) if gs.values[0] is not None else 0
    if values_count != 1 + 26 * object_count:
        return None
    output = GlobalObjectsSample()
    output.time = gs.time
    for i in range(0, object_count):
        obj = DynamicObject()
        b = 1 + 26 * i
        obj.id = int(gs.values[b] if gs.values[b] is not None else 0)
        obj.type = DynamicObjectType(int(gs.values[b + 1])) if gs.values[b + 1] is not None else DynamicObjectType.UNKNOWN
        obj.center_longitude = float(gs.values[b + 2]) if gs.values[b + 2] is not None else 0.0
        obj.center_latitude = float(gs.values[b + 3]) if gs.values[b + 3] is not None else 0.0
        obj.center_altitude = float(gs.values[b + 4]) if gs.values[b + 4] is not None else None
        obj.orientation = float(gs.values[b + 5]) if gs.values[b + 5] is not None else None
        obj.width = float(gs.values[b + 6]) if gs.values[b + 6] is not None else None
        obj.length = float(gs.values[b + 7]) if gs.values[b + 7] is not None else None
        obj.height = float(gs.values[b + 8]) if gs.values[b + 8] is not None else None
        obj.speed = float(gs.values[b + 9]) if gs.values[b + 9] is not None else None
        obj.accel_x = float(gs.values[b + 10]) if gs.values[b + 10] is not None else None
        obj.curvature = float(gs.values[b + 11]) if gs.values[b + 11] is not None else None
        obj.yaw_rate = float(gs.values[b + 12]) if gs.values[b + 12] is not None else None
        obj.pitch = float(gs.values[b + 13]) if gs.values[b + 13] is not None else None
        obj.roll = float(gs.values[b + 14]) if gs.values[b + 14] is not None else None
        obj.rear_axis_longitudinal_offset = float(gs.values[b + 15]) if gs.values[b + 15] is not None else None
        obj.rear_axis_height_offset = float(gs.values[b + 16]) if gs.values[b + 16] is not None else None
        if gs.values[b + 17] is not None and gs.values[b + 18] is not None and gs.values[b + 19] is not None:
            obj.color_valid = True
            obj.color_r = int(gs.values[b + 17])
            obj.color_g = int(gs.values[b + 18])
            obj.color_b = int(gs.values[b + 19])
        output.objects.append(obj)
    return output


def _interpolate_angle(a1: float, w1: float, a2: float, w2: float) -> float | None:
    deg2rad = pi / 180
    x1 = cos(a1 * deg2rad)
    y1 = sin(a1 * deg2rad)
    x2 = cos(a2 * deg2rad)
    y2 = sin(a2 * deg2rad)
    xo = x1 * w1 + x2 * w2
    yo = y1 * w1 + y2 * w2
    if xo == 0 and yo == 0:
        return None
    return atan2(yo, xo) / deg2rad


def _interpolate_global_objects_sample(s1: GlobalObjectsSample, w1: float, s2: GlobalObjectsSample, w2: float) -> GlobalObjectsSample:
    output = GlobalObjectsSample()
    output.time = bi.time
    for o1 in s1.objects:
        id = o1.id
        for o2 in s2.objects:
            if o2.id != id:
                continue
            obj = DynamicObject()
            obj.id = id
            obj.type = o1.type if w1 > w2 else o2.type
            obj.center_longitude = o1.center_longitude * w1 + o2.center_longitude * w2
            obj.center_latitude = o1.center_latitude * w1 + o2.center_latitude * w2
            if o1.center_altitude is not None and o2.center_altitude is not None:
                obj.center_altitude = o1.center_altitude * w1 + o2.center_altitude * w2
            if o1.orientation is not None and o2.orientation is not None:
                obj.orientation = _interpolate_angle(o1.orientation, w1, o2.orientation, w2)
            if o1.width is not None and o2.width is not None:
                obj.width = o1.width * w1 + o2.width * w2
            if o1.length is not None and o2.length is not None:
                obj.length = o1.length * w1 + o2.length * w2
            if o1.height is not None and o2.height is not None:
                obj.height = o1.height * w1 + o2.height * w2
            if o1.speed is not None and o2.speed is not None:
                obj.speed = o1.speed * w1 + o2.speed * w2
            if o1.accel_x is not None and o2.accel_x is not None:
                obj.accel_x = o1.accel_x * w1 + o2.accel_x * w2
            if o1.curvature is not None and o2.curvature is not None:
                obj.curvature = o1.curvature * w1 + o2.curvature * w2
            if o1.yaw_rate is not None and o2.yaw_rate is not None:
                obj.yaw_rate = o1.yaw_rate * w1 + o2.yaw_rate * w2
            if o1.pitch is not None and o2.pitch is not None:
                obj.pitch = o1.pitch * w1 + o2.pitch * w2
            if o1.roll is not None and o2.roll is not None:
                obj.roll = o1.roll * w1 + o2.roll * w2
            if o1.rear_axis_longitudinal_offset is not None and o2.rear_axis_longitudinal_offset is not None:
                obj.rear_axis_longitudinal_offset = o1.rear_axis_longitudinal_offset * w1 + o2.rear_axis_longitudinal_offset * w2
            if o1.rear_axis_height_offset is not None and o2.rear_axis_height_offset is not None:
                obj.rear_axis_height_offset = o1.rear_axis_height_offset * w1 + o2.rear_axis_height_offset * w2
            if o1.color_valid and o2.color_valid:
                obj.color_valid = True
                obj.color_r = o1.color_r if w1 > w2 else o2.color_r
                obj.color_g = o1.color_g if w1 > w2 else o2.color_g
                obj.color_b = o1.color_b if w1 > w2 else o2.color_b
            output.objects.append(obj)
    return output


# Query GlobalObjectsSample for input / 获取GlobalObjectsSample，用于样本输入
def get_global_objects_sample() -> GlobalObjectsSample | None:
    s1 = None
    s2 = None
    w1 = 0.0
    w2 = 0.0
    protocol_id_v1 = 'global-objects-sample-v1'
    protocol_id = 'global-objects-sample-v2'
    if protocol_id in bi.input_samples:
        pair = bi.input_samples[protocol_id]
        s1 = _conv_global_objects_sample_v2(pair.sample1)
        w1 = pair.weight1
        s2 = _conv_global_objects_sample_v2(pair.sample2)
        w2 = pair.weight2
    elif protocol_id_v1 in bi.input_samples:
        pair = bi.input_samples[protocol_id_v1]
        s1 = _conv_global_objects_sample_v1(pair.sample1)
        w1 = pair.weight1
        s2 = _conv_global_objects_sample_v1(pair.sample2)
        w2 = pair.weight2
    if s1 is not None and s2 is not None:
        return _interpolate_global_objects_sample(s1, w1, s2, w2)
    return None
